import os
import dashscope
import requests
import json
from dashscope import Generation
from typing import Dict, Any, TypedDict, List, Tuple
from langgraph.graph import StateGraph, END
from dotenv import load_dotenv


load_dotenv()  # 加载 .env 文件
dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")

# 定义状态结构
class AgentState(TypedDict):
    user_input: str
    intent: str
    intent_details: Dict[str, Any]
    response: str
    knowledge_base_result: Any
    llm_response: str


# 初始化状态
def init_state(user_input: str) -> AgentState:
    return {
        "user_input": user_input,
        "intent": "",
        "intent_details": {},
        "response": "",
        "knowledge_base_result": None,
        "llm_response": ""
    }


# 意图识别智能体
class IntentClassifier:
    def __init__(self):
        # 场景定义知识库
        self.scenes = {
            1: "日常对话",
            2: "新客户发展场景",
            3: "存量客户运营场景",
            4: "客户高价值经营场景",
            5: "混合场景"
        }

        # 场景描述Prompt模板
        self.scene_prompt = """
        请将用户问题分类到以下业务场景：
        1. 日常对话 - 问候、闲聊、非业务相关对话
        2. 新客户发展场景 - 客户获取、市场拓展、异网用户分析等
        3. 存量客户运营场景 - 客户保留、忠诚度分析、全球通用户分析等
        4. 客户高价值经营场景 - ARPU分析、收入增长、高价值客户识别等
        5. 混合场景 - 同时涉及多个业务场景的问题

        用户问题：{user_input}

        请只返回场景编号（1-5），不要包含任何其他内容。
        """

    def classify(self, text: str) -> Tuple[int, str, Dict]:
        """使用通义千问进行意图分类"""
        prompt = self.scene_prompt.format(user_input=text)

        try:
            response = Generation.call(
                model='qwen-max',
                prompt=prompt,
                max_tokens=10,
                temperature=0.1
            )

            # 提取场景编号
            output = response.output.text
            scene_id = int(output) if output.isdigit() else 1

            # 确保在有效范围内
            scene_id = scene_id if 1 <= scene_id <= 5 else 1

            return scene_id, self.scenes[scene_id], {}
        except Exception as e:
            print(f"意图分类失败: {str(e)}")
            return 1, "日常对话", {}


# 通义千问客户端封装
class QwenClient:
    def generate_response(self, prompt: str, system_message: str = "") -> str:
        """使用通义千问生成响应"""
        messages = []
        if system_message:
            messages.append({"role": "system", "content": system_message})
        messages.append({"role": "user", "content": prompt})

        try:
            response = Generation.call(
                model='qwen-max',
                messages=messages,
                max_tokens=2000,
                temperature=0.7
            )
            return response.output.choices[0].message.content
        except Exception as e:
            return f"大模型请求失败: {str(e)}"


# 知识库查询封装
class KnowledgeBaseClient:
    def __init__(self, endpoint: str):
        self.endpoint = endpoint

    def execute_query(self, query: str, params: List[Any] = []) -> Any:
        payload = {
            "context_id": "knowledge_query",
            "query": query,
            "params": params
        }

        try:
            response = requests.post(self.endpoint, json=payload)
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"error": f"知识库查询失败: {str(e)}"}


# 创建智能体节点
def create_intent_classifier_node(classifier: IntentClassifier):
    def classify_intent(state: AgentState) -> AgentState:
        user_input = state["user_input"]
        scenario_id, scenario_name, details = classifier.classify(user_input)

        state["intent"] = scenario_name
        state["intent_details"] = {
            "scenario_id": scenario_id,
            "details": details
        }

        print(f"意图识别结果: {scenario_name} (ID: {scenario_id})")
        return state

    return classify_intent


def create_daily_dialog_node(qwen_client: QwenClient):
    def handle_daily_dialog(state: AgentState) -> AgentState:
        user_input = state["user_input"]

        # 调用通义千问生成响应
        response = qwen_client.generate_response(
            user_input,
            "你是一个专业的客户服务助手，请用友好、专业的语气回答用户的问题。"
        )

        state["response"] = response
        state["llm_response"] = response
        return state

    return handle_daily_dialog


def create_scene_handler_node(scene_id: int, kb_client: KnowledgeBaseClient, qwen_client: QwenClient):
    def handle_scene(state: AgentState) -> AgentState:
        user_input = state["user_input"]

        # 使用通义千问生成SQL查询
        sql_prompt = f"""
        你是一个SQL专家，请将用户问题转换为有效的SQL查询语句。
        数据库表结构：
        - customers (id, name, type, location, arpu, create_date)
        - products (id, name, price, category)

        用户问题：{user_input}

        请只返回SQL语句，不要包含任何解释。
        """

        sql_query = qwen_client.generate_response(sql_prompt)

        # 清理SQL语句
        if "```sql" in sql_query:
            sql_query = sql_query.split("```sql")[1].split("```")[0].strip()
        elif "```" in sql_query:
            sql_query = sql_query.split("```")[1].strip()

        print(f"生成的SQL查询: {sql_query}")

        # 查询知识库
        result = kb_client.execute_query(sql_query)

        if "error" in result:
            state["response"] = f"查询失败: {result['error']}"
            return state

        # 使用通义千问解释结果
        explain_prompt = f"""
        用户问题：{user_input}
        SQL查询：{sql_query}
        查询结果：{json.dumps(result)}

        请完成以下任务：
        1. 用自然语言解释查询结果
        2. 分析数据中的关键洞察
        3. 基于业务场景提供建议
        4. 如果数据量较大，总结主要趋势

        业务场景：{state['intent']}
        """

        response = qwen_client.generate_response(
            explain_prompt,
            "你是一个数据分析专家，请解释查询结果并提供业务建议。"
        )

        state["response"] = response
        state["knowledge_base_result"] = result
        state["llm_response"] = response
        return state

    return handle_scene


# 创建混合场景处理器
def create_mixed_scene_node(kb_client: KnowledgeBaseClient, qwen_client: QwenClient):
    def handle_mixed_scene(state: AgentState) -> AgentState:
        user_input = state["user_input"]

        # 使用通义千问分解问题
        decompose_prompt = f"""
        用户问题包含多个业务场景，请将其分解为2-3个独立的子问题：
        原始问题：{user_input}

        请按以下格式返回：
        [
          {{"sub_question": "子问题1", "scene_type": "场景类型"}},
          {{"sub_question": "子问题2", "scene_type": "场景类型"}}
        ]
        """

        decomposition = qwen_client.generate_response(decompose_prompt)

        try:
            # 尝试解析JSON
            sub_questions = json.loads(decomposition)
        except:
            # 解析失败则创建默认分解
            sub_questions = [{"sub_question": user_input, "scene_type": "混合场景"}]

        # 处理每个子问题
        results = []
        for sub_q in sub_questions:
            # 创建子状态
            sub_state = {
                "user_input": sub_q["sub_question"],
                "intent": sub_q["scene_type"],
                "intent_details": {}
            }

            # 根据场景类型路由处理
            if "新客户" in sub_q["scene_type"]:
                handler = create_scene_handler_node(2, kb_client, qwen_client)
            elif "存量" in sub_q["scene_type"]:
                handler = create_scene_handler_node(3, kb_client, qwen_client)
            elif "高价值" in sub_q["scene_type"]:
                handler = create_scene_handler_node(4, kb_client, qwen_client)
            else:
                handler = create_daily_dialog_node(qwen_client)

            sub_state = handler(sub_state)
            results.append({
                "question": sub_q["sub_question"],
                "answer": sub_state["response"]
            })

        # 综合所有结果
        synthesis_prompt = f"""
        原始问题：{user_input}
        子问题解答：{json.dumps(results)}

        请综合以上信息，提供一个完整的回答：
        1. 总结各个子问题的关键发现
        2. 分析问题之间的关联
        3. 提供整体业务建议
        """

        response = qwen_client.generate_response(synthesis_prompt)
        state["response"] = response
        state["knowledge_base_result"] = results
        state["llm_response"] = response
        return state

    return handle_mixed_scene


# 创建路由函数
def route_based_on_intent(state: AgentState) -> str:
    intent = state["intent"]

    if "日常对话" in intent:
        return "daily_dialog"
    elif "新客户发展" in intent:
        return "new_customer"
    elif "存量客户运营" in intent:
        return "existing_customer"
    elif "客户高价值经营" in intent:
        return "high_value"
    elif "混合场景" in intent:
        return "mixed_scene"
    else:
        return "daily_dialog"  # 默认处理


# 构建智能体工作流
def create_agent_workflow():
    # 初始化组件
    classifier = IntentClassifier()
    qwen_client = QwenClient()
    kb_client = KnowledgeBaseClient(endpoint="http://your-mcp-server/mcp/query")

    # 创建节点
    intent_node = create_intent_classifier_node(classifier)
    daily_dialog_node = create_daily_dialog_node(qwen_client)
    new_customer_node = create_scene_handler_node(2, kb_client, qwen_client)
    existing_customer_node = create_scene_handler_node(3, kb_client, qwen_client)
    high_value_node = create_scene_handler_node(4, kb_client, qwen_client)
    mixed_scene_node = create_mixed_scene_node(kb_client, qwen_client)

    # 构建工作流
    workflow = StateGraph(AgentState)

    # 添加节点
    workflow.add_node("classify_intent", intent_node)
    workflow.add_node("daily_dialog", daily_dialog_node)
    workflow.add_node("new_customer", new_customer_node)
    workflow.add_node("existing_customer", existing_customer_node)
    workflow.add_node("high_value", high_value_node)
    workflow.add_node("mixed_scene", mixed_scene_node)

    # 设置入口点
    workflow.set_entry_point("classify_intent")

    # 添加条件路由
    workflow.add_conditional_edges(
        "classify_intent",
        route_based_on_intent,
        {
            "daily_dialog": "daily_dialog",
            "new_customer": "new_customer",
            "existing_customer": "existing_customer",
            "high_value": "high_value",
            "mixed_scene": "mixed_scene",
        }
    )

    # 设置结束点
    workflow.add_edge("daily_dialog", END)
    workflow.add_edge("new_customer", END)
    workflow.add_edge("existing_customer", END)
    workflow.add_edge("high_value", END)
    workflow.add_edge("mixed_scene", END)

    # 编译工作流
    return workflow.compile()


# 运行智能体
def run_agent(user_input: str):
    # 初始化状态
    state = init_state(user_input)

    # 创建工作流
    agent = create_agent_workflow()

    # 执行工作流
    result = agent.invoke(state)

    return result["response"]


# 测试智能体
if __name__ == "__main__":
    # 测试用例
    test_cases = [
        "你好，最近有什么优惠活动吗？",
        "北京家庭圈中有多少异网用户？",
        "全球通用户在总用户数中的占比是多少？",
        "哪些用户的ARPU有增长的趋势？",
        "北京地区全球通用户的ARPU增长情况如何？"
    ]

    print("通义千问智能体测试结果：")
    print("=" * 60)
    for i, query in enumerate(test_cases, 1):
        print(f"测试 {i}: {query}")
        response = run_agent(query)
        print(f"响应: {response}")
        print("-" * 60)
